import numpy as np
import torch
from gensim.models import FastText
from sklearn.metrics.pairwise import cosine_similarity


# 加载预先训练好的 FastText 模型
class FastTextModel(object):
    def __init__(self):
        self.model = FastText.load_fasttext_format('/home/Dyf/code/models/pretrain_models/wiki.zh/wiki.zh.bin')

    def embedding_text(self, text):
        # 将文本转换为词袋表示，假设我们已经进行了必要的分词和预处理
        words = text.split()
        # 获取每个短语所有单词向量的平均值作为整体的向量表示
        avg_vec = np.mean([self.model.wv[word] for word in words if word in self.model.wv], axis=0)
        return avg_vec
    def embedding_text_to_tensor(self, text):
        # 将文本转换为词袋表示，假设我们已经进行了必要的分词和预处理
        words = text.split()
        # 获取每个短语所有单词向量的平均值作为整体的向量表示
        avg_vec = np.mean([self.model.wv[word] for word in words if word in self.model.wv], axis=0)
        return torch.from_numpy(avg_vec)

    @staticmethod
    def compute_similarity(vec1, vec2):
        phrase_sim = cosine_similarity([vec1], [vec2])[0][0]
        return phrase_sim

    def compute_cosine_similarity(self, text1, text2):
        vec1 = self.embedding_text(text1)
        vec2 = self.embedding_text(text2)
        return self.compute_similarity(vec1, vec2)


FTModel = FastTextModel()
# a = FTModel.compute_cosine_similarity("aaa", "bbbb")
# print(a)
# print(FTModel.embedding_text_to_tensor("aaaaaa"))